from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select
from sqlalchemy import update
from fastapi import HTTPException, status

from app.db.models.credit import Credit, CreditLog
from app.db.models.user import User
from app.schemas.credit import CreditLogCreate

async def get_user_credits(db: AsyncSession, user_id: int) -> Credit:
    """Gets the credit balance for a user. Creates record if not found."""
    result = await db.execute(
        select(Credit).filter(Credit.user_id == user_id)
    )
    credit = result.scalars().first()
    # Create credit record if it doesn't exist for the user
    if not credit:
        credit = await initialize_user_credits(db, user_id)
    return credit

async def initialize_user_credits(db: AsyncSession, user_id: int) -> Credit:
    """Creates the initial credit record for a user (usually 0)."""
    db_credit = Credit(user_id=user_id, balance=0)
    db.add(db_credit)
    await db.commit()
    await db.refresh(db_credit)
    return db_credit

async def add_credit_log(db: AsyncSession, log_entry: CreditLogCreate) -> CreditLog:
    """Adds a record to the credit log."""
    db_log = CreditLog(**log_entry.model_dump())
    db.add(db_log)
    await db.commit()
    await db.refresh(db_log)
    return db_log

async def adjust_credits(db: AsyncSession, user_id: int, amount: int, reason: str | None = None, operator_id: int | None = None) -> Credit:
    """Adjusts user credits and logs the transaction. Handles potential negative balance."""
    # Check if credit record exists, create if not
    credit_record = await get_user_credits(db, user_id)

    # Use UPDATE...WHERE to prevent race conditions
    stmt = (
        update(Credit)
        .where(Credit.user_id == user_id)
        .where(Credit.balance + amount >= 0) # Ensure balance doesn't go negative
        .values(balance=Credit.balance + amount)
        .returning(Credit.balance)
    )
    result = await db.execute(stmt)
    updated_balance = result.scalar_one_or_none()

    if updated_balance is None:
        # The update failed, likely due to insufficient balance
        await db.rollback() # Rollback any potential changes from get_user_credits if it created the record
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Insufficient credits. Current balance: {credit_record.balance}, Amount to deduct: {-amount if amount < 0 else amount}"
        )

    # Log the transaction if successful
    log_entry = CreditLogCreate(
        user_id=user_id,
        amount=amount,
        reason=reason,
        operator_id=operator_id
    )
    await add_credit_log(db, log_entry)

    await db.commit() # Commit both the update and the log
    await db.refresh(credit_record) # Refresh to get the latest state
    return credit_record

async def deduct_credits_for_generation(db: AsyncSession, user_id: int, cost: int = 1) -> Credit:
    """Specific function to deduct credits for image generation."""
    return await adjust_credits(db, user_id, -cost, reason="image_generation")

async def admin_set_credits(db: AsyncSession, user_id: int, new_balance: int, admin_user_id: int) -> Credit:
    """Allows an admin to directly set a user's credit balance."""
    if new_balance < 0:
         raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Credit balance cannot be negative."
        )

    credit_record = await get_user_credits(db, user_id)
    current_balance = credit_record.balance
    amount_changed = new_balance - current_balance

    stmt = (
        update(Credit)
        .where(Credit.user_id == user_id)
        .values(balance=new_balance)
    )
    await db.execute(stmt)

    # Log the adjustment
    log_entry = CreditLogCreate(
        user_id=user_id,
        amount=amount_changed,
        reason="admin_adjust",
        operator_id=admin_user_id
    )
    await add_credit_log(db, log_entry)

    await db.commit()
    await db.refresh(credit_record)
    return credit_record
